#!/usr/bin/env python3
# 数据集中喜剧类型数据最少，可以用来快速体验
import os
import random
import argparse
import logging
import numpy as np
from tensorboardX import SummaryWriter

from libbots import data, model, utils

import torch
import torch.optim as optim
import torch.nn.functional as F

SAVES_DIR = "saves"

BATCH_SIZE = 32
LEARNING_RATE = 1e-3
MAX_EPOCHES = 100

log = logging.getLogger("train")

TEACHER_PROB = 0.5 #设置每次训练使用的teacher forcing方式的概率


def run_test(test_data, net, end_token, device="cpu"):
    '''
    todo 测试训练完成后的网络
    使用加载数据的5%计算每轮epoch的bleu得分
    '''
    bleu_sum = 0.0
    bleu_count = 0
    for p1, p2 in test_data:
        input_seq = model.pack_input(p1, net.emb, device)
        enc = net.encode(input_seq)
        _, tokens = net.decode_chain_argmax(enc, input_seq.data[0:1],
                                            seq_len=data.MAX_TOKENS,
                                            stop_at_token=end_token)
        bleu_sum += utils.calc_bleu(tokens, p2[1:])
        bleu_count += 1
    return bleu_sum / bleu_count


if __name__ == "__main__":
    logging.basicConfig(format="%(asctime)-15s %(levelname)s %(message)s", level=logging.INFO)
    parser = argparse.ArgumentParser()
    # 训练参数解析器
    # --data: 选择使用哪种类型的数据进行训练
    # --cuda: 选择使用哪种设备进行训练
    # --name: 训练的名字
    parser.add_argument("--data", required=True, help="Category to use for training. "
                                                      "Empty string to train on full dataset")
    parser.add_argument("--cuda", action='store_true', default=False,
                        help="Enable cuda")
    parser.add_argument("-n", "--name", required=True, help="Name of the run")
    args = parser.parse_args()
    device = torch.device("cuda" if args.cuda else "cpu")

    # 创建训练存储目录
    saves_path = os.path.join(SAVES_DIR, args.name)
    os.makedirs(saves_path, exist_ok=True)

    # phrase_pairs： 对话对
    # emb_dict： 词嵌入字典
    phrase_pairs, emb_dict = data.load_data(genre_filter=args.data)
    log.info("Obtained %d phrase pairs with %d uniq words",
             len(phrase_pairs), len(emb_dict))
    # todo 这里保存进去的作用是什么
    data.save_emb_dict(saves_path, emb_dict)
    end_token = emb_dict[data.END_TOKEN]
    # 将对话转化为可训练的由token id组成的对白对
    train_data = data.encode_phrase_pairs(phrase_pairs, emb_dict)
    # todo 这个是什么作用
    rand = np.random.RandomState(data.SHUFFLE_SEED)
    # 打乱训练集顺序
    rand.shuffle(train_data)
    log.info("Training data converted, got %d samples", len(train_data))
    train_data, test_data = data.split_train_test(train_data)
    log.info("Train set has %d phrases, test %d", len(train_data), len(test_data))
    # 构建短语模型
    net = model.PhraseModel(emb_size=model.EMBEDDING_DIM, dict_size=len(emb_dict),
                            hid_size=model.HIDDEN_STATE_SIZE).to(device)
    log.info("Model: %s", net)

    # 构建训练统计器
    writer = SummaryWriter(comment="-" + args.name)

    # 构建优化器
    optimiser = optim.Adam(net.parameters(), lr=LEARNING_RATE)
    best_bleu = None
    for epoch in range(MAX_EPOCHES):
        losses = []
        bleu_sum = 0.0
        bleu_count = 0
        # 从train_data中获取BATCH_SIZE个对话
        for batch in data.iterate_batches(train_data, BATCH_SIZE):
            optimiser.zero_grad()
            # 将对白转换为可输入网络的格式
            input_seq, out_seq_list, _, out_idx = model.pack_batch(batch, net.emb, device)
            # 将输入序列编码为中间编码格式
            enc = net.encode(input_seq)

            net_results = [] # 存储网络预测到的回答字符的token id的概率分布序列
            net_targets = [] # 存储label 回答的token id序列（去除了第一个token）
            # 遍历下一句话list
            # 编码是一次性传入整个batch，解码是需要一句一句解码
            for idx, out_seq in enumerate(out_seq_list):
                # idx 索引
                # out_seq 输出的下一句话

                # 提取下一句回答token id序列，并去除第一个token
                ref_indices = out_idx[idx][1:]
                enc_item = net.get_encoded_item(enc, idx)
                # 这边随机值是用来做什么？以下两个计算评分的方式有什么区别
                # 了解decode_teacher就知道，这里是为了加速训练以及防止过拟合
                # 才采用随机一种序列化编码方式，可参考读书笔记joplin
                if random.random() < TEACHER_PROB:
                    # 计算模型根据输入预测出来的下一句话的BLEU值评分
                    # todo 这里返回的r的序列长度是否和另一个评分方法中的一致，感觉这里返回的序列长度是完整的，下一个方法中返回的序列长度和ref_indices一样是少了第一个字符token id
                    r = net.decode_teacher(enc_item, out_seq)
                    bleu_sum += model.seq_bleu(r, ref_indices)
                else:
                    # 得到根据下一句回答字符第一个字符预测到的完整的回答字符的token id概率分布以及token id的序列
                    r, seq = net.decode_chain_argmax(enc_item, out_seq.data[0:1],
                                                     len(ref_indices))
                    # 计算bleu评分
                    bleu_sum += utils.calc_bleu(seq, ref_indices)
                net_results.append(r)
                net_targets.extend(ref_indices)
                bleu_count += 1
            # todo torch cat的实际作用 这里是为了计算交叉熵损失
            results_v = torch.cat(net_results)
            targets_v = torch.LongTensor(net_targets).to(device)
            # 计算预测的和实际的label序列之间的交叉熵
            # todo 这里是计算交叉熵，但是targets_v序列好像不是概率分布
            loss_v = F.cross_entropy(results_v, targets_v)
            loss_v.backward()
            optimiser.step()
            # 上面的代码中，预测的时候都需要传入实际label回答的第一个字符token id，为什么这样做，在网络训练完成后如何使用？

            losses.append(loss_v.item())
        # 计算bleu平均分
        bleu = bleu_sum / bleu_count
        # 使用测试集测试得到测试机的bleu平均分
        bleu_test = run_test(test_data, net, end_token, device)
        log.info("Epoch %d: mean loss %.3f, mean BLEU %.3f, test BLEU %.3f",
                 epoch, np.mean(losses), bleu, bleu_test)
        writer.add_scalar("loss", np.mean(losses), epoch)
        writer.add_scalar("bleu", bleu, epoch)
        writer.add_scalar("bleu_test", bleu_test, epoch)
        if best_bleu is None or best_bleu < bleu_test:
            #仅保存bleu最好的训练参数数据
            if best_bleu is not None:
                out_name = os.path.join(saves_path, "pre_bleu_%.3f_%02d.dat" %
                                        (bleu_test, epoch))
                torch.save(net.state_dict(), out_name)
                log.info("Best BLEU updated %.3f", bleu_test)
            best_bleu = bleu_test

        if epoch % 10 == 0:
            # 保存训练中间结果
            out_name = os.path.join(saves_path, "epoch_%03d_%.3f_%.3f.dat" %
                                    (epoch, bleu, bleu_test))
            torch.save(net.state_dict(), out_name)

    writer.close()
